#!/usr/bin/env python
# -*- coding:utf-8 -*-

from Crypto.Cipher import AES
from binascii import a2b_base64, b2a_base64
import random, string
import json

class EncryptAES():
    def __init__(self,key,iv):
        self.key = key.encode(encoding="utf-8")
        self.iv = iv.encode(encoding="utf-8")
        self.mode = AES.MODE_CBC

    # 加密
    # 这里密钥key 长度必须为16（AES-128）、24（AES-192）、或32（AES-256）
    def encrypt(self, text):
        cryptor = AES.new(self.key, self.mode, self.iv)
        length = 16
        count = len(text)
        if count % length != 0:
            add = length - (count % length)
        else:
            add = 0
        text = text + ('\0' * add)
        ciphertext = cryptor.encrypt(text.encode(encoding="utf-8"))
        # 加密后的字符串转化为16进制字符串 ,当然也可以转换为base64加密的内容，可以使用b2a_base64(self.ciphertext)
        return b2a_base64(ciphertext).decode()

    # 解密后，去掉补足的空格用strip() 去掉
    def decrypt(self, text):
        cryptor = AES.new(self.key, self.mode, self.iv)
        plaintext = cryptor.decrypt(a2b_base64(text))
        plaintext = plaintext.decode().strip('\0')
        plaintext = "".join([plaintext.strip().rsplit("}", 1)[0], "}"]) 
        return plaintext

    @staticmethod
    def get_random_string(length='8'):
        return ''.join(random.sample(string.ascii_letters + string.digits, length))


if __name__ == '__main__':
    # aes = EncryptAES('Kg$.jaG.&h!.#*KB','Rfj.Jy@.67h.Gq*k')  # 初始化密钥
    # _str = '{"id":"asdasd","prie":"adasd178123"}'
    # e_result = aes.encrypt(_str)
    # d_result = aes.decrypt(e_result)
    # print(e_result)
    # print(d_result)
    # print(json.loads(d_result))
    pass